package com.xiashitech.agent.instrumentation.dubbo.inst;

import cn.hutool.json.JSONUtil;
import com.alibaba.dubbo.rpc.Result;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.base.AbstractAttributesExtractor;
import com.xiashitech.agent.instrumentation.base.context.ReentryContextStorage;
import com.xiashitech.agent.instrumentation.base.propagator.W3CReentryBaggagePropagator;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.netty.netty_v4_1.dubbo.DubboHandler;
import com.xiashitech.agent.instrumentation.dubbo.common.inst.DubboRequest;
import com.xiashitech.agent.instrumentation.dubbo.common.inst.reentry.DubboReentryHeadersGetter;
import com.xiashitech.agent.instrumentation.dubbo.common.inst.reentry.DubboReentryHeadersSetter;
import com.xiashitech.agent.instrumentation.utils.MachineAttributesUtil;
import com.xiashitech.interfaces.agent.dto.SystemInfoDTO;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;

public class XSDubboAttributesExtractor extends AbstractAttributesExtractor<DubboRequest, Result> {
    @Override
    public void onXSStart(AttributesBuilder attributesBuilder, Context context, DubboRequest dubboRequest) {
      try {
            Baggage baggage = Java8Bridge.currentForBaggage();
            BaggageBuilder baggageBuilder = Java8Bridge.fromContextForBaggage(context).toBuilder();
            if(!baggage.isEmpty()) {
                baggage.forEach((k,v) -> baggageBuilder.put(k,v.getValue()));
            }
            baggage = baggageBuilder.build();
            attributesBuilder.put("xs.userId",baggage.getEntryValue("xs.userId") == null ? "" : baggage.getEntryValue("xs.userId"));
            attributesBuilder.put("xs.loginName",baggage.getEntryValue("xs.loginName") == null ? "" : baggage.getEntryValue("xs.loginName"));
            attributesBuilder.put("xs.sessionId",baggage.getEntryValue("xs.sessionId") == null ? "" : baggage.getEntryValue("xs.sessionId"));
            attributesBuilder.put("xs.requestDate",baggage.getEntryValue("xs.requestDate") == null ? "" : baggage.getEntryValue("xs.requestDate"));
            attributesBuilder.put("xs.sessionUpdate",baggage.getEntryValue("xs.sessionUpdate") == null ? "" : baggage.getEntryValue("xs.sessionUpdate"));
            attributesBuilder.put("xs.systemId", StartConfig.getParamValueFromKey(StartConfig.systemId));
            attributesBuilder.put("xs.systemReg", SystemInfoDTO.getConfig().getSystemExist());
            attributesBuilder.put("xs.sourceSystemId", baggage.getEntryValue("xsAuthSystemId") == null ? StartConfig.getParamValueFromKey(StartConfig.systemId):baggage.getEntryValue("xsAuthSystemId"));
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
    }

    @Override
    public void onXSEnd(AttributesBuilder attributesBuilder, Context context, DubboRequest dubboRequest, Result result, Throwable throwable) {
      MachineAttributesUtil.fillMachineAttributes(attributesBuilder);
        if(result == null) {
            return;
        }
        try {
            boolean isServer = dubboRequest.getContext().isProviderSide();
           String url = String.valueOf(dubboRequest.getUrl());
            if(isServer) {
               W3CReentryBaggagePropagator.getInstance().inject(Java8Bridge.currentForXSContext(),result, DubboReentryHeadersSetter.INSTANCE);
                System.out.println("###baggage dubbo server "+ Java8Bridge.currentForXSContext().toString());
                System.out.println("###baggage dubbo server result "+ result.toString());
                //清空response中的上下文。
                //ReentryContextStorage.INSTANCE.close();
            } else {
                System.out.println("###baggage dubbo  client "+ Java8Bridge.currentForXSContext().toString());
                W3CReentryBaggagePropagator.getInstance().extract(Java8Bridge.currentForXSContext(), result, DubboReentryHeadersGetter.INSTANCE).makeCurrent();
            }
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
    }

    @Override
    public String getAttributesExtractorName() {
        return XSDubboAttributesExtractor.class.getName();
    }
}
